from mcp.server.auth.provider import (
    OAuthAuthorizationServerProvider,
    AuthorizationParams,
    AuthorizationCode,
    RefreshToken,
    AccessToken
)
from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
from typing import Any
import secrets
import time


class LeohOAuthServerProvider(OAuthAuthorizationServerProvider):
    def __init__(self, login_url: str):
        self.login_url = login_url
        self.clients : dict[str, OAuthClientInformationFull] = {}
        self.state_mapping: dict[str, Any] = {}
        self.auth_codes: dict[str, AuthorizationCode] = {}
        self.tokens: dict[str, AccessToken] = {}

    async def get_client(self, client_id: str) -> OAuthClientInformationFull | None:
        return self.clients.get(client_id)

    async def register_client(self, client_info: OAuthClientInformationFull) -> None:
        self.clients[client_info.client_id] = client_info

    async def authorize(self, client: OAuthClientInformationFull, params: AuthorizationParams) -> str:
        # /authorize并不会验证用户的账号和密码，而是记录用户的一些信息，然后返回一个重定向的URL
        # authorize只是简单记录一些信息
        # 浏览器会自动跳转到重定向的URL上去

        # state是由客户端生成的随机字符串
        state = params.state
        self.state_mapping[state] = {
            "client_id":client.client_id,
            "redirect_uri":params.redirect_uri,
            "state":params.state,
            "code_challenge":params.code_challenge,
            "resource":params.resource,
            "redirect_uri_provided_explicitly":params.redirect_uri_provided_explicitly,
        }
        # 返回重定向的URL
        return f"{self.login_url}?state={state}&client_id={client.client_id}"

    async def load_authorization_code(
            self, client: OAuthClientInformationFull, authorization_code: str
    ) -> AuthorizationCode | None:
        return self.auth_codes.get(authorization_code)

    async def exchange_authorization_code(
            self, client: OAuthClientInformationFull, authorization_code: AuthorizationCode
    ) -> OAuthToken:
        if authorization_code.code not in self.auth_codes:
            raise ValueError("授权码不可用！")
        # 生成access_token
        token = f"mcp_{secrets.token_hex(32)}"
        # 将生成的token存储在服务器，下次进行校验的时候，就可以从中获取，来判断是否存在，以及是否过期
        self.tokens[token] = AccessToken(
            token=token,
            client_id=client.client_id,
            scopes=['user'],
            expires_at=int(time.time()) + 3600,
            resource=authorization_code.resource,
        )
        # 生成完了access_token后，授权码就可以删掉了
        del self.auth_codes[authorization_code.code]
        return OAuthToken(
            access_token=token,
            expires_in=3600,
            scope="user"
        )

    async def load_refresh_token(self, client: OAuthClientInformationFull, refresh_token: str) -> RefreshToken | None:
        raise NotImplementedError("refresh token not implemented")

    async def exchange_refresh_token(
            self,
            client: OAuthClientInformationFull,
            refresh_token: RefreshToken,
            scopes: list[str],
    ) -> OAuthToken:
        raise NotImplementedError("refresh token not implemented")

    async def load_access_token(self, token: str) -> AccessToken | None:
        access_token = self.tokens.get(token)
        # 判断token是否过期
        if access_token and access_token.expires_at < time.time():
            del self.tokens[token]
            return None
        return access_token

    async def revoke_token(
            self,
            token: AccessToken | RefreshToken,
    ) -> None:
        if token in self.tokens:
            del self.tokens[token.token]


    async def create_authorization_code(self, state: str) -> AuthorizationCode:
        # 获取之前执行/authorize时，存储的用户授权数据
        state_data = self.state_mapping.get(state)

        redirect_uri = state_data.get('redirect_uri')
        code_challenge = state_data.get('code_challenge')
        redirect_uri_provided_explicitly = state_data.get('redirect_uri_provided_explicitly')
        client_id = state_data.get('client_id')
        resource = state_data.get('resource')

        # 生成授权码
        code = f"mcp_{secrets.token_hex(16)}"
        auth_code = AuthorizationCode(
            code= code,
            scopes= ["user"],
            expires_at= time.time() + 300,
            client_id= client_id,
            code_challenge= code_challenge,
            redirect_uri= redirect_uri,
            redirect_uri_provided_explicitly= redirect_uri_provided_explicitly,
            resource= resource,
        )
        self.auth_codes[code] = auth_code
        return auth_code